package com.ideaaedi.logback.defender.defender;

import ch.qos.logback.classic.spi.ILoggingEvent;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.JSONValidator;
import com.ideaaedi.logback.defender.constant.DefenderConstant;
import com.ideaaedi.logback.defender.strategy.DefenderStrategy;
import com.ideaaedi.logback.defender.util.DefenderUtil;
import org.apache.commons.lang3.ClassUtils;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.helpers.MessageFormatter;

import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * 针对json的脱敏
 *
 * @author JustryDeng
 * @since 2021/7/21 1:15:15
 */
public class DefaultJsonMessageDefender extends AbstractLogbackMessageDefender<String> {
    
    /**
     * 当string类型字段的字段值是 一个json字符串时,是否对该内容中的敏感信息进行脱敏
     * <p>
     *     当开启hitStringValueIsJson后，DefaultJsonMessageDefender仅对没命中的key的字符串value进行分析;<br/>
     *     对于那些命中了的key,哪怕它的值长json样，也不会试着分析该值，而是直接采用该key对应的脱敏策略直接进行脱敏
     * </p>
     */
    private final boolean hitStringValueIsJson;
    
    public DefaultJsonMessageDefender(Map<String, DefenderStrategy> strategies, List<LogbackMessageDefender> plugins,
                                      boolean hitStringValueIsJson) {
        super(strategies, plugins);
        this.hitStringValueIsJson = hitStringValueIsJson;
    }
    
    @Override
    public void doDesensitize(ILoggingEvent event, StringBuilder buffer, Map<String, DefenderStrategy> strategyMap) {
        String formatStr = event.getMessage();
        Object[] parameters = event.getArgumentArray();
        
        int markerCount = StringUtils.countMatches(formatStr, "{}");
        int argumentCount = parameters == null ? 0 : parameters.length;
        int minCount = Math.min(markerCount, argumentCount);
        if (minCount == 0) {
            String formattedMessage = event.getFormattedMessage();
            // json object
            int jsonObjStartIndex = formattedMessage.indexOf('{');
            int jsonObjEndIndex = formattedMessage.lastIndexOf('}');
            boolean jsonObjLegal = jsonObjStartIndex != -1 && jsonObjEndIndex > jsonObjStartIndex;
            if (!jsonObjLegal) {
                buffer.append(formattedMessage);
                return;
            }
            int jsonStartIndex = jsonObjStartIndex;
            int jsonEndIndex = jsonObjEndIndex;
            
            // json array
            int jsonArrStartIndex = formattedMessage.indexOf('[');
            int jsonArrEndIndex = formattedMessage.lastIndexOf(']');
            boolean jsonArrLegal = jsonArrStartIndex != -1
                                    && jsonArrEndIndex > jsonArrStartIndex
                                    && jsonArrStartIndex < jsonObjStartIndex
                                    && jsonArrEndIndex > jsonObjEndIndex;
            if (jsonArrLegal) {
                jsonStartIndex = jsonArrStartIndex;
                jsonEndIndex = jsonArrEndIndex;
            }
            
            int beforeLength = buffer.length();
            String targetJson = formattedMessage.substring(jsonStartIndex, jsonEndIndex + 1);
            String afterDesensitizationStr = doDesensitization(targetJson, strategyMap);
            buffer.setLength(beforeLength);
            buffer.append(formattedMessage).replace(beforeLength + jsonStartIndex, beforeLength + jsonEndIndex + 1, afterDesensitizationStr);
            return;
        }
        String[] strArray = new String[minCount];
        for (int i = 0; i < minCount; i++) {
            // 考虑到记录日志时，占位符不可能有很多, 一般都<=3个， 所以这里就循环(而不批量)脱敏了;
            strArray[i] = doDesensitization(parameters[i], strategyMap);
        }
        buffer.append(MessageFormatter.arrayFormat(formatStr, strArray).getMessage());
    }

    /**
     * 进行脱敏
     *
     * @param data
     *         要脱敏的数据对象
     * @param fieldsStrategyMap
     *         定位要脱敏数据的Object 以及 对应的脱敏策略。
     *         提示: 如果是json或string脱敏， 此T类型一般为String, 对应要脱敏的字段；
     *         如果是regex脱敏， 此T类型一般为Pattern, 可通过其匹配到要脱敏的数据。
     *         提示: 因为Map是通过key的equals()和hashCode()方法来定位key-value的，
     *         所以使用时，需要根据自己的情况决定对不对T进行equals()和hashCode()重写。
     * @return 脱敏后的字符串
     */
    protected String doDesensitization(Object data, Map<String, DefenderStrategy> fieldsStrategyMap) {
        if (data == null) {
            return DefenderConstant.NULL;
        }
        if (data instanceof CharSequence) {
            return handleJsonString(String.valueOf(data), fieldsStrategyMap);
        }
        boolean doNotNeedHandle = ClassUtils.isPrimitiveOrWrapper(data.getClass()) || data instanceof Number;
        if (doNotNeedHandle) {
            return data.toString();
        }
        Object result;
        if (data instanceof Collection || data instanceof Object[]) {
            result = handleJsonArray(JSONArray.parseArray(JSON.toJSONString(data)), fieldsStrategyMap);
            return JSON.toJSONString(result);
        }
        result = handleJsonObject(JSONArray.parseObject(JSON.toJSONString(data)), fieldsStrategyMap, data.getClass());
        return JSON.toJSONString(result);
    }

    /**
     * 脱敏json字符串
     *
     * @param jsonStr
     *         json字符串
     * @param fieldsStrategyMap
     *         要脱敏的 字段名 以及 对应的脱敏策略
     * @return 脱敏后的数据
     */
    protected String handleJsonString(String jsonStr, Map<String, DefenderStrategy> fieldsStrategyMap) {
        /*
         * 不处理简单json 和非法json
         * <p>
         * 这里认为简单json形如 "1456123"、["xxx1", "xxx2"]之类的
         */
        boolean legalJson = StringUtils.containsAny(jsonStr, '{', '}');
        if (!legalJson) {
            return jsonStr;
        }
        legalJson = JSONValidator.from(jsonStr).validate();
        String tmpJsonStr = null;
        if (!legalJson && hitStringValueIsJson) {
            // 排除转义的影响
            tmpJsonStr = jsonStr.replace("\\\"", "\"");
            legalJson = JSONValidator.from(tmpJsonStr).validate();
        }
        if (!legalJson) {
            return jsonStr;
        }
        /// JSONArray  形如: [{"k1":"v2"},{"k2":"v2"}]
        /// JSONObject 形如: {"account": [{"k1":"v2"}, {"k2":"v2"}]}
        Object obj = JSON.parse(tmpJsonStr == null ? jsonStr : tmpJsonStr);
        Object result;
        if (obj instanceof JSONArray) {
            result = handleJsonArray((JSONArray) obj, fieldsStrategyMap);
        } else {
            result = handleJsonObject((JSONObject) obj, fieldsStrategyMap, obj.getClass());
        }
        if (result == null) {
            return null;
        }
        return tmpJsonStr == null ? JSON.toJSONString(result) : JSON.toJSONString(JSON.toJSONString(result));
    }

    /**
     * 脱敏JSONArray
     *
     * @param jsonArray
     *         要脱敏的JSONArray
     * @param fieldsStrategyMap
     *         要脱敏的 字段名 以及 对应的脱敏策略
     * @return 脱敏后的集合/数组
     */
    protected Object[] handleJsonArray(JSONArray jsonArray, Map<String, DefenderStrategy> fieldsStrategyMap) {
        if (jsonArray == null) {
            return null;
        }
        Object[] arr = new Object[jsonArray.size()];
        arr = jsonArray.toArray(arr);
        Object obj;
        for (int i = 0; i < arr.length; i++) {
            obj = arr[i];
            boolean itemIsCharSequence = obj instanceof CharSequence;
            if (hitStringValueIsJson && itemIsCharSequence) {
                arr[i] = handleJsonString(String.valueOf(obj), fieldsStrategyMap);
                continue;
            }
            // 如果数组元素值是一个 普通的 数字、字符串, 那么不需要处理
            boolean doNotNeedHandle = obj == null || obj.getClass().isPrimitive() || itemIsCharSequence || obj instanceof Number;
            if (doNotNeedHandle) {
                continue;
            }
            // 如果值是 数组/集合，那么调用handleJsonArray方法进行处理
            boolean invokeHandleJsonArray = obj instanceof Collection || obj instanceof Object[];
            Object result;
            if (invokeHandleJsonArray) {
                JSONArray tmpJsonArray = obj instanceof JSONArray ? (JSONArray) obj :
                        JSON.parseArray(JSON.toJSONString(obj));
                result = handleJsonArray(tmpJsonArray, fieldsStrategyMap);
            } else {
                JSONObject tmpJsonObject = obj instanceof JSONObject ? (JSONObject) obj :
                        JSON.parseObject(JSON.toJSONString(obj));
                result = handleJsonObject(tmpJsonObject, fieldsStrategyMap, Object.class);
            }
            arr[i] = result;
        }
        return arr;
    }

    /**
     * 脱敏JSONObject
     *
     * @param jsonObject
     *         要脱敏的JSONObject
     * @param fieldsStrategyMap
     *         要脱敏的 字段名 以及 对应的脱敏策略
     * @param clazz
     *         脱敏对象的类型
     * @return 脱敏后的对象
     */
    protected <T> T handleJsonObject(JSONObject jsonObject, Map<String, DefenderStrategy> fieldsStrategyMap, Class<T> clazz) {
        if (jsonObject == null) {
            return null;
        }
        Set<String> targetFields = fieldsStrategyMap.keySet();
        String entryKey;
        Object entryValue;
        for (Map.Entry<String, Object> entry : jsonObject.entrySet()) {
            entryKey = entry.getKey();
            entryValue = entry.getValue();
            if (entryValue == null) {
                continue;
            }
            // 值是否是字符串
            boolean charSequenceValue = entryValue instanceof CharSequence;
            boolean containKey = targetFields.contains(entryKey);
            // 对字符串值value，当其对应于的key并不匹配脱敏字段，但是开启了hitStringValueIsJson时，试着对该value进行脱敏
            if (hitStringValueIsJson && !containKey && charSequenceValue) {
                entry.setValue(handleJsonString(String.valueOf(entryValue), fieldsStrategyMap));
                continue;
            }
            // 值是否是基本类型、数字
            boolean normalValue = ClassUtils.isPrimitiveOrWrapper(entryValue.getClass())  || entryValue instanceof Number;
            boolean needHandle = containKey || (!normalValue && !charSequenceValue);
            if (!needHandle) {
                continue;
            }
            if (normalValue || charSequenceValue) {
                entry.setValue(DefenderUtil.doReplace(String.valueOf(entryValue), fieldsStrategyMap.get(entryKey)));
                continue;
            }
            // 如果值是 数组/集合，那么调用handleJsonArray方法进行处理
            boolean invokeHandleJsonArray = entryValue instanceof Collection || entryValue instanceof Object[];
            if (invokeHandleJsonArray) {
                JSONArray tmpJsonArray = entryValue instanceof JSONArray ? (JSONArray) entryValue :
                        JSON.parseArray(JSON.toJSONString(entryValue));
                Object result = handleJsonArray(tmpJsonArray, fieldsStrategyMap);
                entry.setValue(result);
            } else {
                JSONObject tmpJsonObject = entryValue instanceof JSONObject ? (JSONObject) entryValue :
                        JSON.parseObject(JSON.toJSONString(entryValue));
                Object result = handleJsonObject(tmpJsonObject, fieldsStrategyMap, Object.class);
                entry.setValue(result);
            }
        }
        return JSON.parseObject(jsonObject.toJSONString(), clazz);
    }
    
    @Override
    public boolean support(ILoggingEvent event) {
        return true;
    }
}
